#include <cstdlib>
#include <cstdio>
#include <cstring>
#include <cstdint>

#include <sys/types.h>
#include <sys/socket.h>
#include <poll.h>
#include <pthread.h>
#include <arpa/inet.h>
#include <fcntl.h>

#include <iostream>
#include <vector>
#include <string>
#include <iterator>

#include <openssl/err.h>

#include "chat.h"
#include "server.h"

// ================================================================================

static Server server;

#define DEBUG(x) std::cout << x << std::endl;

// ================================================================================

static int readAll(SSL* ssl, char* buf, size_t len) {
  while (len) {

    if (SSL_pending(ssl) == 0) {
      pollfd pfd;
    
      BIO_get_fd(SSL_get_rbio(ssl), &pfd.fd);
      pfd.events = POLLIN | POLLERR | POLLHUP;
      if (poll(&pfd, 1, -1) <= 0) {
        return -1;
      }

      if (pfd.revents & (POLLERR | POLLHUP)) {
        return -1;
      }
    }

    int res = SSL_read(ssl, buf, len);
    if (res <= 0) {
      res = SSL_get_error(ssl, res);
      if (res == SSL_ERROR_WANT_READ || res == SSL_ERROR_WANT_WRITE) {
        //return res;
        continue;
      }
      
      if (res != SSL_ERROR_NONE) {
        return -1;
      }
    }

    buf += res;
    len -= res;
  }
  
  return 0;
}

static const char* sslError(SSL* ssl, int ret) {
  return ERR_reason_error_string(SSL_get_error(ssl, ret));
}

// --------------------------------------------------------------------------------

class Parser {
public:
  Parser(const char* msg) 
    : _msg(msg) { }

  void next(const char** s, uint32_t* len) {
    *len = *(uint32_t*)_msg;
    *s = _msg + sizeof(uint32_t);
    _msg = *s + *len;
  }

private:
  const char* _msg;
};

// --------------------------------------------------------------------------------

Client::Client(int sock, SSL_CTX* sslCtx) { 
  _sock = sock;
  _auth = false;
  _deauth = false;
  _ready = false;

  _sslConn = SSL_new(sslCtx);
  if (!_sslConn) {
    std::cout << "Failed to create an SSL connection!" << std::endl;
    return;
  }

  if (SSL_set_fd(_sslConn, sock) != 1) {
    std::cout << "Failed to setup SSL BIO!" << std::endl;
    return;
  }

  _ready = true;

  pthread_create(&_worker, NULL, worker, this);
}

Client::~Client() {
  if (_sslConn) {
    SSL_free(_sslConn);
  }

  close(_sock);
}

void Client::sendPlain(const void* data, uint32_t len) {
  SSL_write(_sslConn, data, len);
}

void Client::send(const void* data, uint32_t len) {
  SSL_write(_sslConn, &len, sizeof(uint32_t));
  SSL_write(_sslConn, data, len);
}

void Client::send(const BinStream& s) {
  uint32_t len = s.size();
  send(s.data(), len);
}

void* Client::worker(void* arg) {
  Client* self = (Client*)arg;
  uint32_t msg_len;
  char* msg;
  int ret;
  struct pollfd pfd;

  DEBUG("New client");

  //

  if ((ret = SSL_accept(self->_sslConn)) != 1) {
    std::cout << "Failed to initiate SSL on the socket: " << SSL_get_error(self->_sslConn, ret) << std::endl; //<< " " << sslError(self->_sslConn, ret) << " " << std::endl;
    goto cleanup;
  }

  fcntl(self->_sock, F_SETFL, O_NONBLOCK);

  pfd.fd = self->_sock;
  pfd.events = POLLIN | POLLERR | POLLHUP;

  for (;;) {
    if (poll(&pfd, 1, 20000) == 0) {
      DEBUG(self->name() << ": timed out waiting for a message --- disconnecting" << std::endl);
      break;
    }

    if (pfd.revents & (POLLERR | POLLHUP)) {
      DEBUG(self->name() << ": network error");
      break;
    }    

    //DEBUG("New message");

    if (readAll(self->_sslConn, (char*)&msg_len, sizeof(msg_len)) < 0) {
      DEBUG(self->name() << ": network error");
      break;
    }
    msg = new char[msg_len + 1];
    if (readAll(self->_sslConn, msg, msg_len) < 0) {
      DEBUG(self->name() << ": network error");

      delete [] msg;
      break;
    }
    
    msg[msg_len] = 0;
    server.handle(self, msg);
    delete [] msg;

    if (self->_deauth) {
      break;
    }
  }

  DEBUG(self->name() << " leaves us");

  cleanup:
  server.remove(self);
  return NULL;
}

// --------------------------------------------------------------------------------

Server::Server() {
  static unsigned char* sessId = (unsigned char*)"1234";

  _initialized = false;

  SSL_library_init();
  SSL_load_error_strings();
  ERR_load_crypto_strings();

  _sslCtx = SSL_CTX_new(SSLv3_server_method());

  if (!_sslCtx) {
    perror("Failed to create an OpenSSL context!");
    return;
  }

  if (SSL_CTX_load_verify_locations(_sslCtx, "cacert.pem", 0) != 1) {
    std::cout << "Failed to load CA certificate!" << std::endl;
    return;
  }

  if (SSL_CTX_use_certificate_chain_file(_sslCtx, "chat-server.csr") != 1) {
    std::cout << "Failed to load the server certificate!" << std::endl;
    return;  
  }

  if (SSL_CTX_use_RSAPrivateKey_file(_sslCtx, "chat-server.pri", SSL_FILETYPE_PEM) != 1) {
    std::cout << "Failed to load the server private key!" << std::endl;
    return;  
  }

  if (SSL_CTX_set_session_id_context(_sslCtx, sessId, 4) != 1) {
    std::cout << "Failed to set SSL session ID!" << std::endl;
    return;
  }

  SSL_CTX_set_verify(_sslCtx, SSL_VERIFY_PEER, NULL);
  SSL_CTX_set_tmp_dh(_sslCtx, get_dh());

  svr_sock = socket(AF_INET, SOCK_STREAM, 0);
  if (svr_sock == -1) {
    perror("Failed to create a TCP socket");
    return;
  }

  bcast_sock = socket(AF_INET, SOCK_DGRAM, 0);
  if (bcast_sock == -1) {
    perror("Failed to create a UDP socket");
    return;
  }  

  int reuseaddr_val = 1;
  setsockopt(svr_sock, SOL_SOCKET, SO_REUSEADDR, &reuseaddr_val, sizeof(int));
  setsockopt(bcast_sock, SOL_SOCKET, SO_REUSEADDR, &reuseaddr_val, sizeof(int));

  sockaddr_in_t svr_host;
  svr_host.sin_family = AF_INET;
  svr_host.sin_port = htons(CHAT_PORT);
  svr_host.sin_addr.s_addr = INADDR_ANY;

  if (bind(svr_sock, (sockaddr_t*)&svr_host, sizeof(svr_host)) == -1) {
    perror("Failed to bind to the TCP socket");
    return;
  }

  if (bind(bcast_sock, (sockaddr_t*)&svr_host, sizeof(svr_host)) == -1) {
    perror("Failed to bind to the TCP socket");
    return;
  }  

  listen(svr_sock, 5);

  pthread_spin_init(&_lock, 0);
  pthread_spin_init(&_msgLock, 0);
  _initialized = true;
}

Server:: ~Server() {
  if (bcast_sock != -1) {
    close(bcast_sock);
  }

  if (svr_sock != -1) {
    close(svr_sock);
  }

  if (_sslCtx) {
    SSL_CTX_free(_sslCtx);
  }

  if (_initialized) {
    pthread_spin_destroy(&_lock);
    pthread_spin_destroy(&_msgLock);
  }
}

bool Server::isOk() const {
  return _initialized;
}

void Server::remove(Client* client) {
  pthread_spin_lock(&_lock);
  _clientNames.erase(client->name());
  _clients.remove(client);  
  delete client;
  pthread_spin_unlock(&_lock);  
}

int Server::mainLoop() {
  if (!isOk()) {
    return 1;
  }

  struct pollfd p[2];
    
  p[0].fd = svr_sock;
  p[0].events = POLLIN;
  p[1].fd = bcast_sock;
  p[1].events = POLLIN;

  for (;;) {
    poll(p, 2, -1);

    if (p[0].revents & POLLIN) {
      // Incoming connection
      
      int sock = accept(svr_sock, NULL, NULL);
      if (sock > 0) {
        timeval timeout;
        timeout.tv_sec = 0;
        timeout.tv_usec = 5000000;
        setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));

        addClient(new Client(sock, _sslCtx));
      }
    }

    if (p[1].revents & POLLIN) {
      // Broadcast
    
      sockaddr_t addr;
      socklen_t addrlen = sizeof(addr);
      char msg[1024];

      int len = recvfrom(bcast_sock, msg, sizeof(msg), 0, &addr, &addrlen);
      if (len > 0) {
        handleBroadcast(msg, &addr);
      }
    }
  }

  return 0;
}

void Server::addClient(Client* client) {
  pthread_spin_lock(&_lock);
  _clients.push_back(client);
  pthread_spin_unlock(&_lock);
}

void Server::lockMsg() {
  pthread_spin_lock(&_msgLock);
}

void Server::unlockMsg() {
  pthread_spin_unlock(&_msgLock);
}

void Server::history(uint32_t id, BinStream& s) {
  Messages::iterator it = _messages.begin();
  std::stringstream ss;

  char last_id[32];
  char nmsg[32];

  std::advance(it, id);

  lockMsg();
  sprintf(last_id, "%u", (uint32_t)_messages.size());
  sprintf(nmsg, "%u", (uint32_t)std::distance(it, _messages.end()));

  s << last_id << nmsg;
  for (; it != _messages.end(); ++it) {
    s << it->c_str();
    ss << it->c_str() << std::endl;
  }
  unlockMsg();

  std::cout << "Last Id: " << last_id << std::endl << "Messsages: " << ss.str() << std::endl;
}

void Server::handle(Client* client, const char* msg) {
  BinStream s;
  uint32_t len;
  const char* cmd;

  Parser p(msg);
  p.next(&cmd, &len);

  if (strncmp(cmd, "HELLO", len) == 0) {
    static const char* options = "";
    static const char* authFail = "n";

    const char* cName;

    p.next(&cName, &len);
    client->setName(std::string(cName, len));
    if (_clientNames.find(client->name()) != _clientNames.end()) {
      client->send(authFail, 1);
      return;
    }
    client->authenticate();
    _clientNames.insert(client->name());


    DEBUG("Hello from " << client->name());

    char last_id[32];
    sprintf(last_id, "%u", (uint32_t)_messages.size());

    s << "y" << last_id << options;
    client->send(s);

  } else if (client->authenticated()) {
    if (strncmp(cmd, "GET", len) == 0) {
      uint32_t id = 0;
      const char* sid;

      p.next(&sid, &len);
      sscanf(sid, "%u", &id);

      DEBUG("Sending history to the client " << client->name() << " since " << id);

      history(id, s);
      client->send(s);

      DEBUG("Sending history to the client " << client->name() << " done");      
    } else if (strncmp(cmd, "HISTORY", len) == 0) {
      history(0, s);
      client->send(s);

    } else if (strncmp(cmd, "BYE", len) == 0) {
      client->deauthenticate();

    } else if (strncmp(cmd, "SEND", len) == 0) {
      const char* cmsg;
      std::stringstream ss;

      p.next(&cmsg, &len);

      ss << client->name() << ": " << cmsg;

      lockMsg();
      _messages.push_back(ss.str());
      unlockMsg();

      DEBUG("User message from " << client->name());

    } else if (strncmp(cmd, "LIST", len) == 0) {
      for (Clients::const_iterator it = _clients.begin();
           it != _clients.end();
           ++it) {
        s << (*it)->name().c_str();
      }

      client->send(s);
    } else {
      DEBUG("Invalid command");
    }
  } else {
    DEBUG("Client not authenticated");
  }
}


void Server::handleBroadcast(const char* msg, const sockaddr_t* addr) {
  static const char* hello = "CHAT-SERVER by Kartashov A., Batalov E., and Sinay P.\n";

  replyBroadcast(hello, strlen(hello), addr);
}

void Server::replyBroadcast(const char* msg, size_t len, const sockaddr_t* addr) {
  sendto(bcast_sock, msg, len, 0, addr, sizeof(sockaddr_in_t));
}

// --------------------------------------------------------------------------------

int main(int argc, char* argv[]) {
  return server.mainLoop();
}
